import numpy as np
from .tac_frame import TacFrame
from typing import Tuple
import matplotlib.pyplot as plt

def plot_vector_field(position : np.ndarray, array : np.ndarray, scale : float = 20, title: str = ""):
    '''
    Plot a 3D vector field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (N, 3) representing the positions of the tac3d marker points.
        array (np.ndarray): A 2D numpy array of shape (N, 3) representing N 3D vectors on each marker.
        scale (float, optional): Scale factor for the vectors. Defaults to 20.0.
    '''

    if position.squeeze().shape != (400,3) and position.squeeze().shape != (20,20,3):
        raise ValueError("position must be of shape (400,3) or (20,20,3)")
    if array.squeeze().shape != (400,3) and array.squeeze().shape != (20,20,3):
        raise ValueError("array must be of shape (400,3) or (20,20,3)")
    
    # 转换为400x3的形状
    position = position.reshape(-1,3)
    array = array.reshape(-1,3)
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 由位置场散点拟合表面形貌并绘制

    X = position[:, 0].reshape(20, 20)
    Y = position[:, 1].reshape(20, 20)
    Z = position[:, 2].reshape(20, 20)
    ax.plot_surface(X, Y, Z, alpha=0.5)
    # 画向量场
    ax.quiver(position[:,0], position[:,1], position[:,2],
              array[:,0], array[:,1], array[:,2],
              length=scale, normalize=False, color='r')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_aspect('equal')
    # reduce the margin of the plot
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    
    plt.title(title)
    plt.show()


def plot_force_field(position: np.ndarray, force: np.ndarray, scale: float = 20.0, title: str = "Force Field"):
    '''
    Plot a 3D force field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (N, 3) representing the positions of the tac3d marker points.
        force (np.ndarray): A 2D numpy array of shape (N, 3) representing N 3D forces on each marker.
        scale (float, optional): Scale factor for the forces. Defaults to 20.0.
    '''
    plot_vector_field(position, force, scale, title)

def plot_displacement_field(position: np.ndarray, displacement: np.ndarray, scale: float = 2.0, title: str = "Displacement Field"):
    '''
    Plot a 3D displacement field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (400, 3) or (20,20,3) representing the positions of the tac3d marker points.
        displacement (np.ndarray): A 2D numpy array of shape (400, 3) or (20,20,3) representing N 3D displacements on each marker.
        scale (float, optional): Scale factor for the displacements. Defaults to 2.0.
    '''
    plot_vector_field(position, displacement, scale, title)

def plot_scalar_field(position: np.ndarray, scalar: np.ndarray, title: str = ""):
    '''
    Plot a 2D scalar field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (400, 3) or (20,20,3) representing the positions of the tac3d marker points.
        scalar (np.ndarray): A 1D numpy array of shape (400,) or (20,20) representing the scalar values at each marker.
    '''
    if position.squeeze().shape != (400,3) and position.squeeze().shape != (20,20,3):
        raise ValueError("position must be of shape (400,3) or (20,20,3)")
    if scalar.squeeze().shape != (400,) and scalar.squeeze().shape != (20,20):
        raise ValueError("scalar must be of shape (400,) or (20,20)")
    
    # 转换为400x3的形状
    position = position.reshape(-1,3)
    
    fig = plt.figure()
    # 绘制平面热力图
    # ax = fig.add_subplot(111, projection='3d')
    X = position[:, 0].reshape(20, 20)
    Y = position[:, 1].reshape(20, 20)
    Z = position[:, 2].reshape(20, 20)
    scalar = scalar.reshape(20, 20)
    # ax.plot_surface(X, Y, scalar, alpha=1, cmap='hot')
    plt.contourf(X, Y, scalar, cmap='hot')
    plt.colorbar()
    plt.title(title)
    plt.show()

def plot_scalar_field_animation(position: np.ndarray, scalar_sequence: np.ndarray, interval: int = 20, title: str = "", saved: bool = False):
    '''
    Plot an animated 2D scalar field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (400, 3) or (20,20,3) representing the positions of the tac3d marker points.
        scalar_sequence (np.ndarray): A 2D numpy array of shape (T, 400) or (T, 20,20) representing the scalar values at each marker over time.
        interval (int, optional): Interval between frames in milliseconds. Defaults to 200.
        saved (bool, optional): Whether to save the animation as a gif file. Defaults to False.
    '''
    if position.squeeze().shape[1:] != (400,3) and position.squeeze().shape[1:] != (20,20,3):
        raise ValueError("position must be of shape (400,3) or (20,20,3)")
    if scalar_sequence.squeeze().shape[1:] != (400,) and scalar_sequence.squeeze().shape[1:] != (20,20):
        raise ValueError("scalar_sequence must be of shape (T, 400) or (T, 20,20)")
    
    if scalar_sequence.shape[0] != position.shape[0]:
        raise ValueError("The first dimension of scalar_sequence must match the first dimension of position")
    
    # 转换为400x3的形状
    length = scalar_sequence.shape[0]
    
    fig = plt.figure()
    scalar_max = np.max(scalar_sequence)
    scalar_min = np.min(scalar_sequence)
    plt.xlabel('X')
    plt.ylabel('Y')

    def update(frame):
        plt.clf()
        X = position[frame, ..., 0].reshape(20, 20)
        Y = position[frame, ..., 1].reshape(20, 20)
        scalar = scalar_sequence[frame].reshape(20, 20)
        plt.contourf(X, Y, scalar, cmap='coolwarm')
        cbar = plt.colorbar()
        # fix the maximum and minimum of colorbar
        plt.clim(scalar_min, scalar_max)
        plt.title(f"{title} - Frame {frame}({frame/30:.2f}s)")

    from matplotlib.animation import FuncAnimation
    ani = FuncAnimation(fig, update, frames=scalar_sequence.shape[0], interval=interval)
    if saved:
        ani.save(f"{title}.gif", writer='pillow', fps=1000/interval)
    plt.show()

def plot_vector_field_animation(position: np.ndarray, vector_sequence: np.ndarray, interval: int = 20, scale: float = 20.0, title: str = "", saved: bool = False):
    '''
    Plot an animated 3D vector field using matplotlib.
    Args:
        position (np.ndarray): A 2D numpy array of shape (400, 3) or (20,20,3) representing the positions of the tac3d marker points.
        vector_sequence (np.ndarray): A 2D numpy array of shape (T, 400, 3) or (T, 20,20,3) representing T sets of N 3D vectors on each marker over time.
        interval (int, optional): Interval between frames in milliseconds. Defaults to 200.
        scale (float, optional): Scale factor for the vectors. Defaults to 20.0.
        saved (bool, optional): Whether to save the animation as a gif file. Defaults to False.
    '''
    if position.squeeze().shape[1:] != (400,3) and position.squeeze().shape[1:] != (20,20,3):
        raise ValueError("position must be of shape (400,3) or (20,20,3)")
    if vector_sequence.squeeze().shape[1:] != (400,3) and vector_sequence.squeeze().shape[1:] != (20,20,3):
        raise ValueError("vector_sequence must be of shape (T, 400, 3) or (T, 20,20,3)")
    
    if vector_sequence.shape[0] != position.shape[0]:
        raise ValueError("The first dimension of vector_sequence must match the first dimension of position")
    
    # 转换为400x3的形状
    length = vector_sequence.shape[0]
    
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')


    def update(frame):
        ax.cla()
        X = position[frame, ..., 0].reshape(20,20)
        Y = position[frame, ..., 1].reshape(20,20)
        Z = position[frame, ..., 2].reshape(20,20)
        U = vector_sequence[frame, ..., 0].reshape(-1)
        V = vector_sequence[frame, ..., 1].reshape(-1)
        W = vector_sequence[frame, ..., 2].reshape(-1)
        ax.plot_surface(X, Y, Z, alpha=0.5)
        # 画向量场
        ax.quiver(position[frame,...,0], position[frame,...,1], position[frame,...,2],
                U, V, W, length=scale, normalize=False, color='r')
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_aspect('equal')

    from matplotlib.animation import FuncAnimation
    ani = FuncAnimation(fig, update, frames=vector_sequence.shape[0], interval=interval)
    if saved:
        ani.save(f"{title}.gif", writer='pillow', fps=1000/interval)
    plt.show()
